# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/5 11:17 
@Author  : ZhangShenao 
@File    : xhs_agent.py 
@Desc    : 小红书营销智能体
"""

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langgraph.graph.state import CompiledStateGraph, StateGraph

from base import BaseAgentState, get_llm


# 定义小红书Agent状态
class XHSAgentState(BaseAgentState):
    """小红书Agent状态"""
    pass


# 定义小红书Agent节点
def xhs_agent_node(state: XHSAgentState, config: RunnableConfig) -> XHSAgentState:
    # 获取LLM
    llm = get_llm()

    # 构造Prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "你是一位小红书文案大师，请根据用户传递的商品名，生成一篇关于该商品的小红书笔记文案，注意风格活泼，多使用emoji表情。"),
        ("human", "{query}"),
    ])

    # 构造Chain
    chain = prompt | llm | StrOutputParser()

    # 执行Chain,获取结果
    reply = chain.invoke({"query": state["query"]})

    # 返回更新状态
    return {"xhs_content": reply}


# 定义创建子图的函数
def create_xhs_graph() -> CompiledStateGraph:
    """创建小红书子图"""

    # 创建图架构
    graph_builder = StateGraph(XHSAgentState)

    # 添加节点
    graph_builder.add_node(node="xhs_agent", action=xhs_agent_node)

    # 添加边
    graph_builder.set_entry_point(key="xhs_agent")  # 设置入口节点
    graph_builder.set_finish_point(key="xhs_agent")  # 设置结束节点

    # 将图编译成Runnable可执行组件,并返回
    return graph_builder.compile()
